package main

import (
	"context"
	"flag"
	"fmt"
	"log"
	"net/http"
	"net/url"
	"os"
	"os/signal"
	"strings"
	"syscall"
	"time"

	"encoding/json"

	"mcp-server/internal"
	_ "mcp-server/proxy/sse"
	_ "mcp-server/registry/nacos"
	_ "mcp-server/servers/gorm"

	"github.com/gin-gonic/gin"
	"github.com/mark3labs/mcp-go/server"
	"gopkg.in/yaml.v3"

	"github.com/gin-contrib/pprof"
)

const (
	DefaultServerName = "go-mcp-server"
	Version           = "v1"
)

type config struct {
	ssePathSuffix                string
	messagePathSuffix            string
	useFullURLForMessageEndpoint bool
	enable                       bool
	matchList                    []internal.MatchRule
	servers                      []*server.SSEServer
	defaultServer                *server.SSEServer
	redisClient                  *internal.RedisClient
}

// StructToMap 将结构体转换为 map[string]any
func StructToMap(s interface{}) (map[string]any, error) {
	// 将结构体转换为 JSON 字节切片
	jsonData, err := json.Marshal(s)
	if err != nil {
		return nil, err
	}
	// 定义一个 map 用于存储解析后的 JSON 数据
	var result map[string]any
	// 将 JSON 字节切片解析为 map
	err = json.Unmarshal(jsonData, &result)
	if err != nil {
		return nil, err
	}
	return result, nil
}

// 从conf/conf.yaml加载配置
func loadConfig() (*config, error) {
	configPath := "conf/conf.yaml"
	data, err := os.ReadFile(configPath)
	if err != nil {
		return nil, fmt.Errorf("failed to read config file: %v", err)
	}

	var yamlConfig struct {
		MCPServer struct {
			SSEPathSuffix       string `yaml:"sse_path_suffix"`
			MessagePathSuffix   string `yaml:"message_path_suffix"`
			FullMessageEndpoint bool   `yaml:"full_message_endpoint"`
			Enable              bool   `yaml:"enable"`
			Redis               struct {
				Address  string `yaml:"address"`
				Username string `yaml:"username"`
				Password string `yaml:"password"`
				DB       int    `yaml:"db"`
			} `yaml:"redis"`
			MatchList []struct {
				MatchRuleDomain string `yaml:"match_rule_domain"`
				MatchRulePath   string `yaml:"match_rule_path"`
				MatchRuleType   string `yaml:"match_rule_type"`
			} `yaml:"match_list"`
			Servers []struct {
				Name   string                 `yaml:"name"`
				Type   string                 `yaml:"type"`
				Path   string                 `yaml:"path"`
				Config map[string]interface{} `yaml:"config"`
			} `yaml:"servers"`
		} `yaml:"mcpServer"`
	}

	if err := yaml.Unmarshal(data, &yamlConfig); err != nil {
		return nil, fmt.Errorf("failed to parse config file: %v", err)
	}

	// 检查是否启用MCP Server
	if !yamlConfig.MCPServer.Enable {
		return nil, fmt.Errorf("MCP Server is disabled in configuration")
	}

	// 创建Redis客户端
	redisClient, err := internal.NewRedisClient(&internal.RedisConfig{
		Address:  yamlConfig.MCPServer.Redis.Address,
		Username: yamlConfig.MCPServer.Redis.Username,
		Password: yamlConfig.MCPServer.Redis.Password,
		DB:       yamlConfig.MCPServer.Redis.DB,
	})
	if err != nil {
		return nil, fmt.Errorf("failed to create Redis client: %v", err)
	}

	// 创建匹配规则
	matchRules := make([]internal.MatchRule, 0, len(yamlConfig.MCPServer.MatchList))
	for _, rule := range yamlConfig.MCPServer.MatchList {
		matchRules = append(matchRules, internal.MatchRule{
			MatchRuleDomain: rule.MatchRuleDomain,
			MatchRulePath:   rule.MatchRulePath,
			MatchRuleType:   internal.RuleType(rule.MatchRuleType),
		})
	}

	cfg := &config{
		ssePathSuffix:                yamlConfig.MCPServer.SSEPathSuffix,
		messagePathSuffix:            yamlConfig.MCPServer.MessagePathSuffix,
		useFullURLForMessageEndpoint: yamlConfig.MCPServer.FullMessageEndpoint,
		enable:                       yamlConfig.MCPServer.Enable,
		matchList:                    matchRules,
		servers:                      []*server.SSEServer{},
		redisClient:                  redisClient,
	}

	// 创建服务器实例
	for _, serverCfg := range yamlConfig.MCPServer.Servers {
		// 根据服务器类型创建不同的服务器实例

		serv := internal.GlobalRegistry.GetServer(serverCfg.Type)

		if serv == nil {
			return nil, fmt.Errorf("server %s is not registered", serverCfg.Type)
		}

		err = serv.ParseConfig(serverCfg.Config)
		if err != nil {
			return nil, fmt.Errorf("failed to parse server config: %w", err)
		}

		serverInstance, err := serv.NewServer(serverCfg.Name)
		if err != nil {
			return nil, fmt.Errorf("failed to initialize DBServer: %w", err)
		}

		cfg.servers = append(cfg.servers, server.NewSSEServer(serverInstance,
			server.WithBasePath(serverCfg.Path),
			server.WithUseFullURLForMessageEndpoint(cfg.useFullURLForMessageEndpoint),
			server.WithSSEEndpoint(cfg.ssePathSuffix),
			server.WithMessageEndpoint(cfg.messagePathSuffix)))
	}

	return cfg, nil
}

// Server 表示一个基于Gin的服务器
type Server struct {
	config   *config
	stopChan chan struct{}
}

// NewServer 创建一个新的服务器实例
func NewServer(cfg *config) *Server {
	return &Server{
		config:   cfg,
		stopChan: make(chan struct{}),
	}
}

// RequestURL 表示一个请求URL
type RequestURL struct {
	method    string
	scheme    string
	host      string
	path      string
	baseURL   string
	parsedURL *url.URL
}

// NewRequestURL 从Gin上下文创建一个新的RequestURL
func NewRequestURL(c *gin.Context) *RequestURL {
	method := c.Request.Method
	scheme := "http"
	if c.Request.TLS != nil {
		scheme = "https"
	}
	host := c.Request.Host
	path := c.Request.URL.Path
	baseURL := fmt.Sprintf("%s://%s", scheme, host)
	parsedURL := c.Request.URL

	log.Printf("RequestURL: method=%s, scheme=%s, host=%s, path=%s\n", method, scheme, host, path)
	return &RequestURL{method: method, scheme: scheme, host: host, path: path, baseURL: baseURL, parsedURL: parsedURL}
}

// HandleSSERequest 处理SSE连接请求
func (s *Server) HandleSSERequest(c *gin.Context) {
	url := NewRequestURL(c)
	path := url.parsedURL.Path
	log.Printf("RequestURL: %+v", url)

	// 检查请求是否匹配match_list中的任何规则
	if !internal.IsMatch(s.config.matchList, url.host, path) {
		fmt.Printf("Request does not match any rule in match_list: %s\n", url.parsedURL.String())
		c.Status(http.StatusNotFound)
		return
	}

	for _, serv := range s.config.servers {
		if path == serv.CompleteSsePath() {
			serv.ServeHTTP(c.Writer, c.Request)
			return
		}
	}

	c.Status(http.StatusNotFound)
}

// HandleMessageRequest 处理消息请求
func (s *Server) HandleMessageRequest(c *gin.Context) {
	url := NewRequestURL(c)
	path := url.parsedURL.Path

	// 检查请求是否匹配match_list中的任何规则
	if !internal.IsMatch(s.config.matchList, url.host, path) {
		fmt.Printf("Request does not match any rule in match_list: %s\n", url.parsedURL.String())
		c.Status(http.StatusNotFound)
		return
	}

	for _, serv := range s.config.servers {
		if path == serv.CompleteMessagePath() {
			if url.method != http.MethodPost {
				c.String(http.StatusMethodNotAllowed, "Method not allowed")
				return
			}

			serv.ServeHTTP(c.Writer, c.Request)
			return
		}
	}

	c.Status(http.StatusNotFound)
}

// Shutdown 关闭服务器
func (s *Server) Shutdown() {
	select {
	case <-s.stopChan:
		return
	default:
		fmt.Println("Stopping SSE connections")
		close(s.stopChan)
	}
}

// setupRouter 设置Gin路由
func setupRouter(server *Server) *gin.Engine {
	router := gin.Default()

	// 添加中间件
	router.Use(gin.Recovery())
	router.Use(gin.Logger())

	pprof.Register(router)

	// mcp base path + sse path
	router.Any("/:base/*sse", func(c *gin.Context) {
		path := c.Request.URL.Path

		if has := strings.HasSuffix(path, server.config.ssePathSuffix); has {
			if server.config.defaultServer != nil && path == server.config.defaultServer.CompleteSsePath() {
				server.HandleSSERequest(c)
				return
			}

			for _, s := range server.config.servers {
				log.Println("request path:", path, s.CompleteSsePath())
				if path == s.CompleteSsePath() {
					server.HandleSSERequest(c)
					return
				}
			}

			return
		}

		if has := strings.HasSuffix(path, server.config.messagePathSuffix); has {
			for _, s := range server.config.servers {
				log.Println("request path:", path, s.CompleteMessagePath())

				if server.config.defaultServer != nil && path == server.config.defaultServer.CompleteSsePath() {
					server.HandleSSERequest(c)
					return
				}

				if path == s.CompleteMessagePath() {
					server.HandleMessageRequest(c)
					return
				}
			}
		}

		c.Status(http.StatusNotFound)
	})

	return router
}

func main() {
	// 解析命令行参数
	var transport string
	port := flag.Int("port", 8080, "Server port")
	flag.StringVar(&transport, "t", "sse", "Transport type (stdio or sse)")
	flag.StringVar(&transport, "transport", "sse", "Transport type (stdio or sse)")
	flag.Parse()

	log.SetFlags(log.Lshortfile | log.LstdFlags)

	// 加载配置
	cfg, err := loadConfig()
	if err != nil {
		log.Fatalf("Failed to load configuration: %v", err)
	}

	// 创建服务器实例
	server := NewServer(cfg)

	// 设置Gin路由
	router := setupRouter(server)

	// 创建HTTP服务器
	srv := &http.Server{
		Addr:    fmt.Sprintf(":%d", *port),
		Handler: router,
	}

	// 优雅关闭
	go func() {
		// 等待中断信号
		quit := make(chan os.Signal, 1)
		signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM)
		<-quit
		log.Println("Shutting down server...")

		// 关闭SSE连接
		server.Shutdown()

		// 设置5秒的超时时间关闭服务器
		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer cancel()
		if err := srv.Shutdown(ctx); err != nil {
			log.Fatal("Server forced to shutdown:", err)
		}

		log.Println("Server exiting")
	}()

	// 启动服务器
	log.Printf("Server is running on port %d", *port)
	if err := srv.ListenAndServe(); err != nil && err != http.ErrServerClosed {
		log.Fatalf("Failed to start server: %v", err)
	}
}
